package pdftable;


import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.rendering.ImageType;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripperByArea;
import org.apache.pdfbox.tools.imageio.ImageIOUtil;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.apache.poi.ss.usermodel.Cell;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.ss.util.CellRangeAddress;
import org.opencv.core.Core;
import org.opencv.core.Point;
import org.opencv.core.Rect;
import pdftable.models.ParsedTablePage;

import java.awt.*;
import java.awt.geom.Rectangle2D.Double;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeSet;
import java.util.stream.Collectors;

import static pdftable.Utils.bufferedImage2GrayscaleMat;

@Slf4j
public class PdfTableReader {

    private TableExtractor extractor;
    private PdfTableSettings settings;

    static {
        System.loadLibrary(Core.NATIVE_LIBRARY_NAME);
    }

    public PdfTableReader(PdfTableSettings settings) {
        this.settings = settings;
        this.extractor = new TableExtractor(settings);
    }

    public PdfTableReader() {
        this(new PdfTableSettings());
    }


    public List<ParsedTablePage> pdfTable2Excel(PDDocument document, int startPage, int endPage, OutputStream out) throws Exception {
        List<ParsedTablePage> tablePages = parsePdfTablePages(document, startPage, endPage);
        builderExcel(tablePages,out);
        return tablePages;
    }
    static void builderExcel(List<ParsedTablePage> tablePages, OutputStream out) throws Exception {
        if(log.isDebugEnabled()){
            log.debug("tablePages={}",JSON.toJSONString(tablePages));
        }
        ArrayList<TreeSet> treeSetXs = new ArrayList<TreeSet>();
        ArrayList<TreeSet> treeSetYs = new ArrayList<TreeSet>();
        for (ParsedTablePage tablePage : tablePages) {
            TreeSet<Integer> xs = new TreeSet<Integer>();
            TreeSet<Integer> ys = new TreeSet<Integer>();
            List<ParsedTablePage.ParsedTableRow> rows = tablePage.getRows();
            for (ParsedTablePage.ParsedTableRow row : rows) {
                List<Rect> rects = row.getRects();
                for (int i = 0; i < rects.size(); i++) {
                    Rect rect = rects.get(i);
                    xs.add(rect.x);
                    ys.add(rect.y);
                }
            }
            treeSetXs.add(xs);
            treeSetYs.add(ys);
        }
        Workbook wb = new HSSFWorkbook();
        int sheet_i = 0;
        for (ParsedTablePage tablePage : tablePages) {
            Sheet sheet = wb.createSheet("sheet"+sheet_i);
            TreeSet<Integer> xs = treeSetXs.get(sheet_i);
            TreeSet<Integer> ys = treeSetYs.get(sheet_i);
            createTable(xs,ys,sheet);
            List<ParsedTablePage.ParsedTableRow> rows = tablePage.getRows();
            int row_i = 0;
            for (ParsedTablePage.ParsedTableRow row : rows) {
                List<String> cells = row.getCells();
                List<Rect> rects = row.getRects();
                for (int i = 0; i < rects.size(); i++) {
                    Rect rect = rects.get(i);
                    setVal(xs,ys,sheet,rect,cells.get(i));
                }

            }
            sheet_i++;
        }
        wb.write(out);
        wb.close();
    }
    static void setVal(TreeSet<Integer> txs,TreeSet<Integer> tys,Sheet sheet,Rect rect,String val){
        ArrayList<Integer> xs = new ArrayList<>(txs);
        ArrayList<Integer> ys = new ArrayList<>(tys);
        int x = rect.x;
        int y = rect.y;
        int maxX = rect.width + x;
        int maxY = rect.height + y;
        int firstRow = ys.indexOf(y);
        int firstCol = xs.indexOf(x);

        int lastRow =firstRow;
        int lastCol =firstCol;
        for (int i = 0; i < xs.size(); i++) {
            if(Math.abs(xs.get(i)-maxX)<=2){
                lastCol=i-1;
                break;
            }
        }
        for (int i = 0; i < ys.size(); i++) {
            if(Math.abs(ys.get(i)-maxY)<=2){
                lastRow=i-1;
                break;
            }
        }


        Row row = sheet.getRow(firstRow);
        Cell cell = row.getCell(firstCol);
        cell.setCellValue(val);
        if(firstCol<lastCol||firstRow<lastRow){
            try {
                sheet.addMergedRegion(new CellRangeAddress(
                        firstRow, //first row (0-based)
                        lastRow, //last row  (0-based)
                        firstCol, //first column (0-based)
                        lastCol  //last column  (0-based)
                ));
            }catch (Exception e){
                e.printStackTrace();
                throw e;
            }

        }
    }
    static void createTable(TreeSet xs,TreeSet ys,Sheet sheet){
        for (int i = 0; i < ys.size(); i++) {
            Row row = sheet.createRow( i);
            for (int j = 0; j < xs.size(); j++) {
                row.createCell(j);
            }
        }
    }
    /**
     * Renders PDF page with DPI specified in settings and saves it in specified directory.
     *
     * @param renderer  PDF renderer instance
     * @param page      page number
     * @param outputDir output directory
     * @throws IOException
     */
    private void savePdfPageAsPNG(PDFRenderer renderer, int page, Path outputDir) throws IOException {
        BufferedImage bim;
        synchronized (this) {
            bim = renderer.renderImageWithDPI(page, settings.getPdfRenderingDpi(), ImageType.RGB);
        }
        Path outPath = outputDir.resolve(Paths.get("page_" + (page + 1) + ".png"));
        ImageIOUtil.writeImage(bim, outPath.toString(), settings.getPdfRenderingDpi());

    }

    /**
     * Renders PDF pages range with DPI specified in settings and saves images in specified directory.
     *
     * @param document  PDF document instance
     * @param startPage first page in range (first page == 1)
     * @param endPage   last page in range
     * @param outputDir output directory
     * @throws IOException
     */
    public void savePdfPagesAsPNG(PDDocument document, int startPage, int endPage, Path outputDir) throws IOException {
        PDFRenderer pdfRenderer = new PDFRenderer(document);
        for (int page = startPage - 1; page < endPage; ++page) {
            savePdfPageAsPNG(pdfRenderer, page, outputDir);
        }
    }

    /**
     * Renders single PDF page with DPI specified in settings and saves image in specified directory.
     *
     * @param document  PDF document instance
     * @param page      page number (first page == 1)
     * @param outputDir output directory
     * @throws IOException
     */
    public void savePdfPageAsPNG(PDDocument document, int page, Path outputDir) throws IOException {
        savePdfPagesAsPNG(document, page, page, outputDir);
    }

    /**
     * Parses single PDF page and returns list of rows containing cell texts.
     *
     * @param bi     PDF page in image format
     * @param pdPage PDF page in PDPage format
     * @return parsed page
     * @throws IOException
     */
    private ParsedTablePage parsePdfTablePage(BufferedImage bi, PDPage pdPage, int pageNumber) throws IOException {

        Map tableBoundingRectangles = extractor.getTableBoundingRectangles(bufferedImage2GrayscaleMat(bi));
        List<Rect> rectangles =(List<Rect>) tableBoundingRectangles.get("out");
        List<List<Point>> tables =(List<List<Point>>) tableBoundingRectangles.get("tables");
        return parsePageByRectangles(pdPage, rectangles, pageNumber,tables);
    }

    /**
     * Parses range of PDF pages and returns list of lists of rows containing cell texts.
     *
     * @param document  PDF document instance
     * @param startPage first page in range to parse (first page == 1)
     * @param endPage   last page in range
     * @return List of pages
     * @throws IOException
     */
    public List<ParsedTablePage> parsePdfTablePages(PDDocument document, int startPage, int endPage) throws IOException {
        List<ParsedTablePage> out = new ArrayList<>();
        PDFRenderer renderer = new PDFRenderer(document);
        for (int page = startPage - 1; page < endPage; ++page) {
            BufferedImage bi;
            synchronized (this) {
                bi = renderer.renderImageWithDPI(page, settings.getPdfRenderingDpi(), ImageType.RGB);
            }
            ParsedTablePage parsedTablePage = parsePdfTablePage(bi, document.getPage(page), page + 1);
            if(log.isDebugEnabled())
            log.debug("parsedTablePage = {}",JSON.toJSONString(parsedTablePage));
            out.add(parsedTablePage);
        }
        return out;
    }

    /**
     * Parses single PDF page and returns list of rows containing cell texts.
     *
     * @param document PDF document instance
     * @param page     number of page to parse (first page == 1)
     * @return parsed page
     * @throws IOException
     */
    public ParsedTablePage parsePdfTablePage(PDDocument document, int page) throws IOException {
        return parsePdfTablePages(document, page, page).get(0);
    }

    /**
     * Saves debug images of PDF pages from specified range and saves them in specified directory.
     *
     * @param document  PDF document instance
     * @param startPage first page in range to process (first page == 1)
     * @param endPage   last page in range
     * @param outputDir destination directory
     * @throws IOException
     */
    public void savePdfTablePagesDebugImages(PDDocument document, int startPage, int endPage, Path outputDir) throws IOException {
        TableExtractor debugExtractor = new TableExtractor(settings);
        PDFRenderer renderer = new PDFRenderer(document);
        for (int page = startPage - 1; page < endPage; ++page) {
            PdfTableSettings debugSettings = PdfTableSettings.getBuilder()
                    .setDebugImages(true)
                    .setDebugFileOutputDir(outputDir)
                    .setDebugFilename("page_" + (page + 1))
                    .build();
            debugExtractor.setSettings(debugSettings);
            BufferedImage bi;
            synchronized (this) {
                bi = renderer.renderImageWithDPI(page, settings.getPdfRenderingDpi(), ImageType.RGB);
            }
            debugExtractor.getTableBoundingRectangles(bufferedImage2GrayscaleMat(bi));
        }
    }

    /**
     * Saves debug images of PDF page and saves them in specified directory.
     *
     * @param document  PDF document instance
     * @param page      page to process (first page == 1)
     * @param outputDir destination directory
     * @throws IOException
     */
    public void savePdfTablePageDebugImage(PDDocument document, int page, Path outputDir) throws IOException {
        savePdfTablePagesDebugImages(document, page, page, outputDir);
    }

    /**
     * Parses PDF page cell by cell using rectangles obtained from TableExtractor.
     *
     * @param page       PDF page
     * @param rectangles list of OpenCV rectangles recognized by TableExtractor
     * @return parsed page
     * @throws IOException
     */
    private ParsedTablePage parsePageByRectangles(PDPage page, List<Rect> rectangles, int pageNumber,List<List<Point>> tables) throws IOException {
        List<List<Rect>> sortedRects = groupRectanglesByRow(rectangles,tables);
        ParsedTablePage out = new ParsedTablePage(pageNumber);

        PDFTextStripperByArea stripper = new PDFTextStripperByArea();
        stripper.setSortByPosition(true);

        int iRow = 0;
        int iCol = 0;
        for (List<Rect> row : sortedRects) {
            for (Rect col : row) {
                Rectangle r = new Rectangle(
                        (int) (col.x * settings.getDpiRatio()),
                        (int) (col.y * settings.getDpiRatio()),
                        (int) (col.width * settings.getDpiRatio()),
                        (int) (col.height * settings.getDpiRatio())
                );
                Double s = new Double( col.x * settings.getDpiRatio(),  col.y * settings.getDpiRatio(), col.width * settings.getDpiRatio() , col.height * settings.getDpiRatio());
                stripper.addRegion(getRegionId(iRow, iCol), s);
                iCol++;
            }
            iRow++;
            iCol = 0;
        }
        if(log.isDebugEnabled()){
            log.debug("List<Rect> = {}", JSON.toJSONString(sortedRects));
        }
        stripper.extractRegions(page);

        iRow = 0;
        iCol = 0;
        for (List<Rect> row : sortedRects) {
            List<String> rowCells = new ArrayList<>();
            for (Rect col : row) {
                String cellText = stripper.getTextForRegion(getRegionId(iRow, iCol));

                if(cellText==null ||"\r\n".equals(cellText)){
                     cellText = "  ";
                }else {
                    cellText =  cellText.replaceAll("\\s", "");
                }
                rowCells.add(cellText);
                iCol++;
            }

            out.addRow(rowCells);
            ParsedTablePage.ParsedTableRow parsedTableRow =
                    out.getRows().get(out.getRows().size() - 1);
            parsedTableRow.setRects(row);
            iRow++;
            iCol = 0;
        }

        return out;
    }


//    String aa2222(Rect col,PDPage page){
//        PDFTextStripperByArea stripper = null;
//        try {
//            stripper = new PDFTextStripperByArea();
//        } catch (IOException e) {
//            e.printStackTrace();
//        }
//        stripper.setSortByPosition(true);
//        Rectangle r = new Rectangle(
//                (int) (col.x * settings.getDpiRatio()),
//                (int) (col.y * settings.getDpiRatio()),
//                (int) (col.width * settings.getDpiRatio()),
//                (int) (col.height * settings.getDpiRatio()+1)
//        );
//        stripper.addRegion(getRegionId(1, 1), r);
//        try {
//            stripper.extractRegions(page);
//            return stripper.getTextForRegion(getRegionId(1, 1));
//        } catch (IOException e) {
//            e.printStackTrace();
//        }
//        return "";
//    }
    /**
     * Groups rectangles by y coordinate effectively grouping them into rows.
     *
     * @param rectangles list of OpenCV Rectangles
     * @return list of Rectangle lists representing table rows.
     */
    private List<List<Rect>> groupRectanglesByRow(List<Rect> rectangles,List<List<Point>> tables) {
        List<List<Rect>> out = new ArrayList<>();
        List<List<Rect>> out2 = new ArrayList<>();
        if(tables.isEmpty()) return out;
        List<Integer> rowsCoords = rectangles.stream().map(r -> r.y).distinct().collect(Collectors.toList());
        for (int rowCoords : rowsCoords) {
            List<Rect> cols = rectangles.stream().filter(r -> r.y == rowCoords).collect(Collectors.toList());
            out.add(cols);
        }
        for (List<Point> table : tables) {
            Point point1 = table.get(0);
            Point point2 = table.get(1);
            Point point3 = table.get(2);
            Point point4 = table.get(3);

            //double minx  = point1.x>point2.x?point2.x:point1.x;
            for (int i = 0; i < out.size(); i++) {
                Rect rect = out.get(i).get(0);
                if(rect.y>=point1.y && rect.y<=point3.y){
                    out2.add(out.get(i));
                }
            }

        }



        return out2;
    }

    /**
     * Static helper for creating row/column markers.
     *
     * @param row table row
     * @param col table column
     * @return marker with row & column number
     */
    private static String getRegionId(int row, int col) {
        return String.format("r%dc%d", row, col);
    }

}
